import os
import json
from PIL import Image
import shutil
from tqdm import tqdm
import subprocess

def simple_crop(image_path, save_root, crop_w_num, crop_h_num, resize=None, image_name=None):
    """
    对图像进行规则切割，并保存切割后的区域

    Args:
        image_path (str): 原始图像路径
        save_root (str): 保存根目录
        crop_w_num (int): 水平方向切割数量
        crop_h_num (int): 垂直方向切割数量
        resize (tuple or int, optional): 调整大小的目标尺寸，如果是int则表示正方形
        image_name (str, optional): 图像名称，如果未提供则从image_path提取

    Returns:
        dict: 包含保存目录信息的字典
    """
    # 如果未提供image_name，从image_path中提取文件名（不含扩展名）
    if not image_name:
        image_name = os.path.basename(image_path).split('.')[0]
    
    # 创建保存目录
    save_dir = os.path.join(save_root, image_name)
    os.makedirs(save_dir, exist_ok=True)
    
    # 使用shutil安全删除目录下所有文件，处理包含空格的文件名
    if os.path.exists(save_dir):
        for file in os.listdir(save_dir):
            file_path = os.path.join(save_dir, file)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f"警告：删除文件 {file_path} 时发生错误: {e}")

    # 打开原始图像
    try:
        img = Image.open(image_path)
    except Exception as e:
        print(f"错误：无法打开图像 {image_path}: {e}")
        return {'save_dir': save_dir}

    # 获取原始图像尺寸
    img_w, img_h = img.size

    # 如果指定了resize，先调整图像大小
    if resize is not None:
        if isinstance(resize, int):
            resize = (resize, resize)
        try:
            img = img.resize(resize, Image.Resampling.LANCZOS)
            img_w, img_h = resize
        except Exception as e:
            print(f"警告：调整图像大小失败 {image_path}: {e}")

    # 计算每个切割区域的宽度和高度
    crop_w = img_w // crop_w_num
    crop_h = img_h // crop_h_num

    # 进行切割
    for i in range(crop_w_num):
        for j in range(crop_h_num):
            # 计算左上角和右下角坐标
            left = i * crop_w
            upper = j * crop_h
            right = min((i + 1) * crop_w, img_w)
            lower = min((j + 1) * crop_h, img_h)

            # 计算左上角点的百分比位置
            percent_x = int((left / img_w) * 100)
            percent_y = int((upper / img_h) * 100)

            # 裁剪图像
            crop_img = img.crop((left, upper, right, lower))

            # 构造保存文件名
            crop_name = f"crop{percent_x}_{percent_y}.jpg"
            save_path = os.path.join(save_dir, crop_name)

            # 保存裁剪后的图像
            try:
                crop_img.save(save_path, quality=95)
            except Exception as e:
                print(f"警告：保存裁剪图像失败 {save_path}: {e}")

    return {
        'save_dir': save_dir,
        'crop_w_num': crop_w_num,
        'crop_h_num': crop_h_num,
        'resize': resize
    }

if __name__ == '__main__':
    # 定义数据路径和参数
    data_path = '/mnt/c/DocumentWorkspace/public_processed'
    original_root = '/mnt/d/pubilc_dataset'
    save_root = './experiments'
    max_crop_num = 50
    record_json = {}

    # 读取annotations.json
    try:
        with open(os.path.join(data_path, 'annotations.json')) as f:
            data_dict = json.load(f)
    except Exception as e:
        print(f"错误：无法读取annotations.json: {e}")
        exit(1)

    # 计算处理总数
    num_total_num = len(data_dict.keys())
    total = max(max_crop_num, num_total_num)

    # 创建保存根目录
    os.makedirs(save_root, exist_ok=True)

    # 处理每张图像
    for image_name in tqdm(data_dict, total=total, unit='images'):
        data = data_dict[image_name]
        original_path = os.path.join(original_root, data['original_path'])

        # 进行切割
        save_info = simple_crop(
            image_path=original_path,
            save_root=os.path.join(save_root, 'simple_crop'),
            crop_w_num=4,
            crop_h_num=4,
            resize=(224, 224),
            image_name=image_name
        )
        record_json[image_name] = save_info

    # 保存记录文件
    try:
        with open(os.path.join(save_root, 'region_record.json'), 'w') as f:
            json.dump(record_json, f, indent=4)
    except Exception as e:
        print(f"错误：无法保存region_record.json: {e}")